execfile('v4/base.py')
import pandas as pd
import numpy as np
from typing import Optional


class TurnoverF(Base):
    def __init__(self):
        self.bar_count_default = 50
        self.end_date_str = '20250831'  # 如果需要可以取消注释
        # self.start_date_str = '2025-06-16'  # 如果需要可以取消注释

    def mark_col_growth(
        self,
        df: pd.DataFrame,
        target_col: str,
        window: int = 10,
        feature_suffix: Optional[str] = None,
        min_valid_value: float = 1e-6,
        growth_type: str = 'relative',
        trading_col: str = 'is_trading',
        is_drop: bool = True
    ) -> pd.DataFrame:
        """
        优化版增长率计算 - 保持原有4列输出
        
        改进点：
        1. 零值特殊处理：当基准值为零时，使用替代策略计算增长率
        2. 仅保留最终需要的4列输出
        3. 使用临时DataFrame避免污染原始数据
        """
        # 输入校验
        required_cols = {'code', 'date', target_col}
        missing_cols = required_cols - set(df.columns)
        if missing_cols:
            raise ValueError(f"缺失必要列: {missing_cols}")
        
        # 创建临时工作DataFrame（仅包含必要列）
        work_df = df.copy()
        
        # 数据预处理
        work_df['date'] = pd.to_datetime(work_df['date'])
        work_df.sort_values(by=['code', 'date'], inplace=True)
        
        # 特征列命名
        suffix = f"_{feature_suffix}" if feature_suffix else ''
        growth_col = f"{target_col}_yoy{suffix}_{window}d"
        pre_growth_col = f"pre_{growth_col}"
        
        # 零值特殊处理 - 添加微小常数避免除零错误
        epsilon = 1e-10
        work_df['_adjusted_target'] = work_df[target_col].replace(0, epsilon) + epsilon
        
        # 滚动基准计算（优化零值处理）
        work_df['_base'] = work_df.groupby('code', group_keys=False)['_adjusted_target'].transform(
            lambda x: x.shift(1).rolling(window, min_periods=1).mean()
        )
        # display(work_df[-30:])
        # 增长率计算（优化零值处理）
        prev_col = work_df.groupby('code')[target_col].shift(1)
        
        if growth_type == 'relative':
            # 核心优化：零值特殊处理
            with np.errstate(divide='ignore', invalid='ignore'):
                work_df[growth_col] = np.where(
                    work_df['_base'] <= min_valid_value,
                    # 基准接近零时的替代计算
                    (work_df[target_col] - prev_col) / (prev_col + epsilon),
                    # 正常相对增长率计算
                    (work_df[target_col] - work_df['_base']) / work_df['_base']
                )
            # display(work_df[-30:])
        elif growth_type == 'log':
            # 对数增长率同样需要零值处理
            work_df[growth_col] = np.where(
                (work_df['_base'] > min_valid_value) & (work_df[target_col] > min_valid_value),
                np.log(work_df[target_col] / work_df['_base']),
                np.where(
                    # 零到非零的跃迁
                    (work_df[target_col] > min_valid_value) & (prev_col <= min_valid_value),
                    1.0,  # 固定值表示显著增长
                    np.nan
                )
            )
        else:
            raise ValueError("growth_type 必须为 'relative' 或 'log'")
        
        # 计算昨日特征
        work_df[pre_growth_col] = work_df.groupby('code')[growth_col].shift(1)
        work_df[growth_col+'_2'] = work_df['_adjusted_target']/work_df['_base']
        
        # 仅保留最终需要的4列
        result = work_df[['code', 'date', growth_col,growth_col+'_2', pre_growth_col]].copy()

        # 停牌数据处理（如果提供停牌列）
        if trading_col in df.columns:
            print('dddd',trading_col)
            # 合并停牌信息
            trading_info = df[['code', 'date', trading_col]].drop_duplicates()
            result = result.merge(trading_info, on=['code', 'date'], how='left')
            result.loc[~result[trading_col], [growth_col, pre_growth_col]] = np.nan
            result.drop(columns=[trading_col], inplace=True)
        
        return result.round(4)  # 提高精度保留小数
    
        
   
    
    def mark_vol_turnover(self,df):
        """生成预测三日涨超15%的核心特征"""
        df = df.sort_values(['code','date']).copy()
        
        # === 特征工程 ===
      
        # 2. 成交额突破信号（网页3突破模型）
        df['turnover_break'] = (df['turnover'] > df.groupby('code')['turnover']
                                .transform(lambda x: x.rolling(5).max().shift())).astype(int)
        
        # 3. 换手率波动强度（网页7标准化处理）
        df['t_rate_zscore'] = df.groupby('code')['t_rate'].transform(
            lambda x: (x - x.rolling(10).mean()) / x.rolling(10).std().replace(0, 1e-6))
        
        return df[['code','date','turnover_break',
                  't_rate_zscore']]